package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"syscall"
)

// humanizeBytes 将字节数转换为人类可读的字符串 (e.g., 1.5 GB)
func humanizeBytes(bytes uint64) string {
	const unit = 1024
	if bytes < unit {
		return fmt.Sprintf("%d B", bytes)
	}
	div, exp := uint64(unit), 0
	for n := bytes / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.2f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
}

// getDiskUsage 获取指定路径所在磁盘分区的使用情况
func getDiskUsage(path string) (total, used, free uint64, err error) {
	var stat syscall.Statfs_t
	err = syscall.Statfs(path, &stat)
	if err != nil {
		return 0, 0, 0, err
	}

	total = stat.Blocks * uint64(stat.Bsize)
	free = stat.Bavail * uint64(stat.Bsize)
	used = total - (stat.Bfree * uint64(stat.Bsize))

	return total, used, free, nil
}

// getDirectorySize 递归计算指定目录的总大小（忽略文件不存在的错误）
func getDirectorySize(path string) (uint64, error) {
	var totalSize uint64
	err := filepath.Walk(path, func(filePath string, info os.FileInfo, err error) error {
		// --- 关键修改1：处理遍历中的错误，忽略“文件不存在”和“权限不足” ---
		if err != nil {
			// 忽略“文件/目录不存在”的错误（如动态删除的文件）
			if os.IsNotExist(err) {
				log.Printf("警告: 路径 '%s' 不存在（已忽略）: %v\n", filePath, err)
				return nil // 返回 nil 继续遍历其他路径，不中断
			}
			// 忽略“权限不足”的错误
			if os.IsPermission(err) {
				log.Printf("警告: 访问 '%s' 权限不足（已忽略）: %v\n", filePath, err)
				return filepath.SkipDir // 跳过无权限的目录
			}
			// 其他错误（如IO错误）返回，避免影响整体遍历
			log.Printf("错误: 处理 '%s' 时出错: %v\n", filePath, err)
			return nil // 也可改为 return err 中断，根据需求选择
		}

		if !info.IsDir() {
			totalSize += uint64(info.Size())
		}
		return nil
	})
	return totalSize, err
}

// getSubdirSizesConcurrently 并发计算指定路径下每个直接子目录的大小（忽略错误）
func getSubdirSizesConcurrently(path string) (map[string]uint64, error) {
	subdirSizes := make(map[string]uint64)
	var wg sync.WaitGroup
	var mu sync.Mutex

	entries, err := os.ReadDir(path)
	if err != nil {
		return nil, err
	}

	for _, entry := range entries {
		if entry.IsDir() {
			wg.Add(1)
			subdirName := entry.Name()
			subdirPath := filepath.Join(path, subdirName)

			go func(name, p string) {
				defer wg.Done()
				size, err := getDirectorySize(p)
				if err != nil {
					log.Printf("警告: 计算子目录 '%s' 大小失败（已忽略）: %v\n", p, err)
					return
				}
				mu.Lock()
				subdirSizes[name] = size
				mu.Unlock()
			}(subdirName, subdirPath)
		}
	}

	wg.Wait()
	return subdirSizes, nil
}

// getTopFiles 获取最大的 N 个文件（忽略文件不存在的错误）
func getTopFiles(path string, topN int) (uint64, []struct {
	path string
	size uint64
}, error) {
	var totalSize uint64
	fileSizes := make(map[string]uint64)

	err := filepath.Walk(path, func(filePath string, info os.FileInfo, err error) error {
		// --- 关键修改2：与 getDirectorySize 一致，忽略文件不存在的错误 ---
		if err != nil {
			if os.IsNotExist(err) {
				log.Printf("警告: 路径 '%s' 不存在（已忽略）: %v\n", filePath, err)
				return nil
			}
			if os.IsPermission(err) {
				log.Printf("警告: 访问 '%s' 权限不足（已忽略）: %v\n", filePath, err)
				return filepath.SkipDir
			}
			log.Printf("错误: 处理 '%s' 时出错: %v\n", filePath, err)
			return nil
		}

		if !info.IsDir() {
			size := uint64(info.Size())
			totalSize += size
			relPath, err := filepath.Rel(path, filePath)
			if err != nil {
				relPath = filePath
			}
			fileSizes[relPath] = size
		}
		return nil
	})

	if err != nil {
		// --- 关键修改3：即使有部分错误，也不直接退出，返回已统计的数据 ---
		log.Printf("警告: 遍历目录时存在部分错误，但已返回可用数据: %v\n", err)
		// 这里不返回 err，而是继续处理已收集的文件大小
	}

	// 转换为切片并排序
	var files []struct {
		path string
		size uint64
	}
	for path, size := range fileSizes {
		files = append(files, struct {
			path string
			size uint64
		}{path: path, size: size})
	}

	sort.Slice(files, func(i, j int) bool {
		return files[i].size > files[j].size
	})

	// 限制返回数量
	if len(files) > topN {
		files = files[:topN]
	}

	return totalSize, files, nil
}

// truncateString 截断过长的字符串（可根据需求调整 maxLen）
func truncateString(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen-3] + "..."
}

func main() {
	// 命令行参数：-n 显示数量，-len 文件名最大显示长度
	topN := flag.Int("n", 10, "指定要显示的最大文件/子目录数量，默认10")
	maxLen := flag.Int("len", 60, "指定文件名/目录名最大显示长度，超过截断，默认60")
	flag.Parse()

	// 获取目标路径
	var targetPath string
	if flag.NArg() > 0 {
		targetPath = flag.Arg(0)
	} else {
		var err error
		targetPath, err = os.Getwd()
		if err != nil {
			log.Fatalf("无法获取当前工作目录: %v", err)
		}
	}

	// 检查路径是否存在（初始路径不存在才退出）
	if _, err := os.Stat(targetPath); os.IsNotExist(err) {
		log.Fatalf("错误: 目标路径 '%s' 不存在", targetPath)
	}

	fmt.Printf("--- 正在分析路径: %s ---\n\n", targetPath)

	// 1. 磁盘分区信息
	fmt.Println("=== 磁盘分区空间占用情况 ===")
	total, used, free, err := getDiskUsage(targetPath)
	if err != nil {
		log.Printf("警告: 无法获取磁盘信息（已忽略）: %v\n", err)
	} else {
		fmt.Printf("  总容量: %s\n", humanizeBytes(total))
		fmt.Printf("  已用空间: %s (%.2f%%)\n", humanizeBytes(used), float64(used)/float64(total)*100)
		fmt.Printf("  可用空间: %s\n", humanizeBytes(free))
	}
	fmt.Println()

	// 2. 目录总大小和Top N文件（关键：即使有错误也继续）
	fmt.Printf("=== 目录内容统计 (Top %d) ===\n", *topN)
	totalSize, topFiles, err := getTopFiles(targetPath, *topN)
	// 不再因 err 退出，而是打印警告
	if err != nil {
		log.Printf("警告: 统计文件时存在部分错误，但已显示可用结果\n")
	}
	fmt.Printf("  目标路径总大小: %s\n", humanizeBytes(totalSize))
	fmt.Printf("  已统计文件数量: %d\n", len(topFiles))
	fmt.Println()

	// 3. 打印Top N文件
	fmt.Printf("=== 最大的 %d 个文件 ===\n", *topN)
	for _, file := range topFiles {
		fmt.Printf("  %-*s %s\n", *maxLen, truncateString(file.path, *maxLen), humanizeBytes(file.size))
	}
	fmt.Println()

	// 4. 并发统计子目录大小
	fmt.Printf("=== 直接子目录大小 (Top %d) ===\n", *topN)
	subdirSizes, err := getSubdirSizesConcurrently(targetPath)
	if err != nil {
		log.Printf("警告: 统计子目录时存在部分错误（已忽略）: %v\n", err)
	} else {
		var subdirs []struct {
			path string
			size uint64
		}
		for dir, size := range subdirSizes {
			subdirs = append(subdirs, struct {
				path string
				size uint64
			}{path: dir, size: size})
		}
		sort.Slice(subdirs, func(i, j int) bool {
			return subdirs[i].size > subdirs[j].size
		})
		if len(subdirs) > *topN {
			subdirs = subdirs[:*topN]
		}
		for _, dir := range subdirs {
			fmt.Printf("  %-*s %s\n", *maxLen, truncateString(dir.path, *maxLen), humanizeBytes(dir.size))
		}
	}

	fmt.Println("\n--- 分析完成（部分忽略错误已在日志中显示）---")
}
